package tsj.ai.alibaba.data.analysis.util;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.document.Document;
import tsj.ai.alibaba.data.analysis.common.util.JsonUtil;

import java.util.HashMap;
import java.util.List;
import java.util.function.Supplier;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-13 15:04:34
 */
public class StateUtil {
    private static final ObjectMapper OBJECT_MAPPER = JsonUtil.getObjectMapper();

    /**
     * Safely get string type state value
     */
    public static String getStringValue(OverAllState state, String key) {
        return state.value(key)
                .map(String.class::cast)
                .orElseThrow(() -> new IllegalStateException("State key not found: " + key));
    }

    /**
     * Safely get string type state value with default value
     */
    public static String getStringValue(OverAllState state, String key, String defaultValue) {
        return state.value(key).map(String.class::cast).orElse(defaultValue);
    }

    /**
     * Safely get list type state value
     */
    @SuppressWarnings("unchecked")
    public static <T> List<T> getListValue(OverAllState state, String key) {
        return state.value(key)
                .map(v -> (List<T>) v)
                .orElseThrow(() -> new IllegalStateException("State key not found: " + key));
    }

    /**
     * Safely get object type state value
     */
    public static <T> T getObjectValue(OverAllState state, String key, Class<T> type) {
        return state.value(key)
                .map(value -> deserializeIfNeeded(value, type))
                .orElseThrow(() -> new IllegalStateException("State key not found: " + key));
    }

    /**
     * Safely get object type state value with default value
     */
    public static <T> T getObjectValue(OverAllState state, String key, Class<T> type, T defaultValue) {
        return state.value(key).map(value -> deserializeIfNeeded(value, type)).orElse(defaultValue);
    }

    /**
     * Handle deserialization of HashMap to target type when needed
     */
    private static <T> T deserializeIfNeeded(Object value, Class<T> type) {
        // If already the correct type, return as-is
        if (type.isInstance(value)) {
            return type.cast(value);
        }

        // If it's a HashMap but we need a complex object, use JSON conversion
        if (value instanceof HashMap && !type.equals(HashMap.class)) {
            return OBJECT_MAPPER.convertValue(value, type);
        }

        return type.cast(value);
    }

    /**
     * Safely get object type state value with default value supplier
     */
    public static <T> T getObjectValue(OverAllState state, String key, Class<T> type, Supplier<T> defaultSupplier) {
        return state.value(key).map(type::cast).orElseGet(defaultSupplier);
    }

    /**
     * Check if state value exists
     */
    public static boolean hasValue(OverAllState state, String key) {
        return state.value(key).isPresent() && !((String) state.value(key).get()).isEmpty();
    }

    /**
     * Get Document list
     */
    public static List<Document> getDocumentList(OverAllState state, String key) {
        return getListValue(state, key);
    }

    /**
     * Get list of Document lists
     */
    public static List<List<Document>> getDocumentListList(OverAllState state, String key) {
        return getListValue(state, key);
    }
}
